Review: PR #699 — feat: make async task cap configurable

Summary

Promotes the previously hard-coded async scheduler submission cap to a public RunConfig.max_submitted_tasks field (default 1024, validated >= 1) and raises the default from 256 to 1024. The async dataset builder now reads the value from run_config and forwards it to AsyncTaskScheduler; max_model_task_admission is computed as max(DEFAULT_TASK_POOL_SIZE, max_submitted_tasks, 2 * aggregate_max_parallel_requests) so model-heavy runs are not capped by the old 256 floor. Capacity reporting now sources submission_capacity from "run_config" rather than "dataset_builder". Total: +48/-8 lines across 7 files.

Findings

Correctness — looks right

• DEFAULT_TASK_POOL_SIZE in async_scheduler.py:80 is now sourced from DEFAULT_SUBMISSION_CAPACITY in task_admission.py:39, giving the two constants a single source of truth. Both default to 1024, matching the new RunConfig default — invariant holds.
• The max_model_task_admission formula at dataset_builder.py:1028-1032 preserves the existing floor semantics: model task admission is at least DEFAULT_TASK_POOL_SIZE even if the user sets max_submitted_tasks below 1024. This matches the PR description's intent.
• Field validation (ge=1) is consistent with neighboring fields like non_inference_max_parallel_workers.
• submission_capacity source string changed to "run_config" accurately reflects the new wiring; downstream consumers of AsyncCapacityPlan are unlikely to depend on the literal string, but worth verifying if any dashboards parse it.

Behavior change worth calling out

The default jumps from 256 → 1024 (4×). Users who never set this knob will see significantly more concurrent task leases. The PR notes this as intended (so model-heavy runs don't get throttled by the old floor), but consider:

• Memory pressure for users with many columns / large row groups: each in-flight task carries scheduler bookkeeping plus any per-task buffers.
• Existing performance baselines or capacity-plan snapshots in tests/recipes may shift; the PR updates the tests that assert submission_capacity directly, but external consumers (logs, telemetry dashboards, runbooks) could be affected.

Not blocking, but a release note / changelog entry beyond the PR body would help downstream users notice.

Naming / minor consistency

• DEFAULT_TASK_POOL_SIZE and DEFAULT_SUBMISSION_CAPACITY are now aliases. Keeping both is fine for layering reasons (the engine module re-exports the config-layer constant), but a one-line comment at async_scheduler.py:80 explaining "alias retained for legacy import surface" would prevent a future refactor from quietly dropping it.
• The docstring in run_config.py:136-137 reads "Maximum number of scheduler task leases, and therefore async scheduler worker tasks…" — the "and therefore" is a bit awkward. Consider: "Maximum number of in-flight scheduler task leases (i.e., the cap on concurrent async scheduler worker tasks)."

Test coverage — solid

• test_run_config.py: default value, custom value, validation rejection — all three cases covered.
• test_async_builder_integration.py: confirms the value flows through _setup_async_pipeline to AsyncTaskScheduler.__init__ (1536 → both max_submitted_tasks and max_model_task_admission, since 1536 > 2 * aggregate in the test).
• test_data_designer.py::test_run_config_setting_persists: round-trip through DataDesigner.set_run_config.
• test_dataset_builder.py: default assertion alongside other run-config defaults.

One gap: there's no test that exercises the floor behavior — i.e., a user setting max_submitted_tasks below DEFAULT_TASK_POOL_SIZE (say 64) and confirming max_model_task_admission stays at 1024. The integration test uses 1536, which exercises the upper path. A small unit test at the builder level locking in the floor would prevent a future refactor from silently dropping it.

Style / project conventions

• Absolute imports preserved.
• from __future__ import annotations present in changed source files.
• Field uses Field(default=…, ge=…, description=…) matching the surrounding style.
• No reverse imports introduced; the new constant lives in the engine package and is consumed there + read out of run_config (config → engine direction is preserved).

Security / performance

• No security-sensitive surface touched.
• Performance: the 4× default raise is the main concern (covered above). Validation is bounded only on the lower side (ge=1); there is no upper bound. A user could set max_submitted_tasks=1_000_000 and OOM. This matches existing fields like buffer_size (also unbounded above), so it's consistent with the project's posture, but a soft sanity ceiling or warning log on unusually large values would be a small defensive improvement.

Verdict

Approve-with-nits. The wiring is clean, tests cover the happy path and validation, and the layering respects the config → engine direction. Suggested non-blocking follow-ups:

1. Add a unit test that pins the max_model_task_admission floor when max_submitted_tasks is set below DEFAULT_TASK_POOL_SIZE.
2. Tighten the max_submitted_tasks docstring and add a one-line comment on the DEFAULT_TASK_POOL_SIZE/DEFAULT_SUBMISSION_CAPACITY alias.
3. Consider a changelog/release-note line flagging the default bump from 256 → 1024 for users with capacity-sensitive deployments.

Greptile Summary

This PR exposes RunConfig.max_in_flight_tasks as a user-facing configuration for the async scheduler's task-lease capacity, raises the default from 256 to 1024, and simplifies internal plumbing by removing the old aggregate-request-headroom multiplier and duplicate constant.

• RunConfig gains a max_in_flight_tasks: int = Field(default=1024, ge=1) field; dataset_builder.py reads it directly and passes it as both max_in_flight_tasks and max_model_task_admission to AsyncTaskScheduler, replacing the prior max(DEFAULT_TASK_POOL_SIZE, 2 × aggregate_max_parallel_requests) computation.
• DEFAULT_IN_FLIGHT_TASK_CAPACITY = 1024 is consolidated in task_admission.py and imported into async_scheduler.py; DEFAULT_TASK_POOL_SIZE and MODEL_TASK_ADMISSION_HEADROOM_MULTIPLIER are deleted.
• The adaptive row-group queue-guard formula is simplified from max(max_submitted_tasks × 4, max_model_task_admission × 2) to max_in_flight_tasks × 4, and the "local" key is removed from the TaskAdmissionConfig resource-limits map (it was never requested by any task type, so this is a no-op for admission control).

Important Files Changed

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    RC[RunConfig\nmax_in_flight_tasks=1024] -->|read| DB[DatasetBuilder\n_prepare_async_run]
    DB -->|max_in_flight_tasks\nmax_model_task_admission| ATS[AsyncTaskScheduler.__init__]
    ATS -->|submission_capacity=max_in_flight_tasks\nresource_limits llm_wait=max_model_task_admission| TAC[TaskAdmissionConfig]
    ATS -->|model_group_limit_cap=max_model_task_admission| TSR[TaskSchedulingResolver]
    ATS -->|_max_in_flight_tasks| QG[Queue Guard\nmax_in_flight_tasks × 4]
    TAC --> CTRL[TaskAdmissionController\nleases submission + llm_wait]
    TSR -->|admitted_limit per model group| CTRL
    QG -->|blocks new row-group admission| RGA[Adaptive Row-Group Admission]

_{Reviews (5): Last reviewed commit: "feat: make async in-flight task cap conf..." | Re-trigger Greptile}